#pragma once

#include "kiva/node.h"

namespace kiva {

class DslNode;

// For any class that wants the full-range expression API, use this.
template <typename Derived, typename T = DslNode>
class ExprAPI {
  // node_type is a value type that is used as arguments and return value.
  using node_type = T;

  using node_initializer = std::initializer_list<node_type>;

  static node_type FromNodeRef(NodeRef node_ref) { return node_ref; }

  static NodeRef ToNodeRef(const node_type& node) { return node; }

  // Requirement to Derived:
  NodeRef ThisToNodeRef() const {
    return static_cast<const Derived*>(this)->ThisToNodeRef();
  }

 public:
  // operator= must be defined by subclass.
  //   node_type operator=(const node_type& node) const {
  // return NewAssign(*this, node);
  //   }
  //   node_type operator=(node_type&& node) const {
  // return NewAssign(*this, node);
  //   }
  node_type operator()() const { return operator()({}); }
  node_type operator()(const node_type& args) const {
    return FromNodeRef(Node::NewCall(ThisToNodeRef(), ToNodeRef(args)));
  }
  // Build a subscript node.
  node_type operator[](const node_type& index) const {
    return FromNodeRef(Node::NewSubscript(ThisToNodeRef(), ToNodeRef(index)));
  }
  // *args or **kwds.
  node_type operator*() const {
    return FromNodeRef(Node::NewUnaryOp(NodeType::kUStar, ThisToNodeRef()));
  }
  // Used for decorator, i.e., &property
  node_type operator&() const {
    return FromNodeRef(
        Node::NewUnaryOp(NodeType::kUAmpersand, ThisToNodeRef()));
  }
  node_type in(const node_type& item) const {
    return FromNodeRef(
        Node::NewCompare(ThisToNodeRef(), NodeType::kIn, ToNodeRef(item)));
  }
  node_type not_in(const node_type& item) const {
    return FromNodeRef(
        Node::NewCompare(ThisToNodeRef(), NodeType::kNotIn, ToNodeRef(item)));
  }
  node_type is(const node_type& item) const {
    return FromNodeRef(
        Node::NewCompare(ThisToNodeRef(), NodeType::kIs, ToNodeRef(item)));
  }
  node_type is_not(const node_type& item) const {
    return FromNodeRef(
        Node::NewCompare(ThisToNodeRef(), NodeType::kIsNot, ToNodeRef(item)));
  }
  node_type as(const char* alias) const {
    return FromNodeRef(Node::NewAlias(ThisToNodeRef(), alias));
  }
  node_type attr(const char* name) const {
    return FromNodeRef(Node::NewAttribute(ThisToNodeRef(), name));
  }
  node_type matmult(const node_type& operand) const {
    return FromNodeRef(Node::NewBinOp(ThisToNodeRef(), NodeType::kMatMult,
                                      ToNodeRef(operand)));
  }
  node_type pow(const node_type& operand) const {
    return FromNodeRef(
        Node::NewBinOp(ThisToNodeRef(), NodeType::kPow, ToNodeRef(operand)));
  }

 protected:
  node_type NewAssign(const node_type& rhs) const {
    return FromNodeRef(Node::NewAssign(ThisToNodeRef(), ToNodeRef(rhs)));
  }
};

class DslNode final : public ExprAPI<DslNode> {
 public:
  // Build Constant Node.
  template <typename T, KIVA_ENABLE_IF(std::is_integral<T>{})>
  DslNode(T val)
      : DslNode(std::is_signed<T>{}
                    ? Node::NewConstant(static_cast<long>(val))
                    : Node::NewConstant(static_cast<unsigned long>(val))) {}

  template <typename T, KIVA_ENABLE_IF(std::is_floating_point<T>{})>
  DslNode(T val) : DslNode(Node::NewConstant(val)) {}

  template <typename T,
            KIVA_ENABLE_IF(std::is_convertible<T, std::string_view>{})>
  DslNode(const T& val) : DslNode(Node::NewConstant(val)) {}

  DslNode(bool val) : DslNode(Node::NewConstant(val)) {}

  using InitializerList = std::initializer_list<DslNode>;

  // Build a list of things into a Tuple.
  DslNode(InitializerList items);
  DslNode operator=(const DslNode& rhs) const;

  DslNode(NodeRef node) : node_(node) {}
  operator NodeRef() const { return node_; }

 private:
  friend class ExprAPI<DslNode>;
  friend NodeRef NodeRefFromDslNode2(const DslNode);
  friend DslNode DslNode2FromNodeRef(NodeRef);

  static NodeRef ToNodeRef(const DslNode& node) { return node.node_; }
  static DslNode FromNodeRef(NodeRef node_ref) { return DslNode(node_ref); }
  NodeRef ThisToNodeRef() const { return node_; }
  static NodeRef NewFromInitializer(InitializerList items);

  NodeRef node_;
};

extern template class ExprAPI<DslNode>;

inline std::unique_ptr<Node> MakeNode(DslNode node) {
  return std::unique_ptr<Node>(NodeRef(node).get());
}

class NameDslNode final : public ExprAPI<NameDslNode> {
 public:
  constexpr NameDslNode(std::string_view val) : value_(val) {}

  operator DslNode() const { return ThisToNodeRef(); }

  DslNode operator=(const DslNode& rhs) const { return NewAssign(rhs); }
  DslNode operator=(const NameDslNode& rhs) const {
    return NewAssign(DslNode(rhs));
  }

 private:
  friend class ExprAPI<NameDslNode>;
  NodeRef ThisToNodeRef() const { return Node::NewName(value_); }
  std::string_view value_;
};

extern template class ExprAPI<NameDslNode>;

// This is the base class for all statement node builders.
class DslNodeBuilder : public NodeBuilder {
 public:
  // Convert to a DslNode as a finalization process.
  operator DslNode();

 protected:
  using NodeBuilder::NodeBuilder;
  DslNodeBuilder(const DslNodeBuilder&) = delete;

  void AddChild(DslNode node);
  void AddChildren(DslNode::InitializerList nodes);
  static DslNode WithType(DslNode node, NodeType new_type);
  static DslNode MakeBody(DslNode node, NodeType type = NodeType::kBody);
};

// Use this macro to give alias and declare an extern template class.
#define KIVA_EXTERN_TEMPLATE_ALIAS(Class, Enum, Name) \
  extern template class Class<NodeType::Enum>;        \
  struct Name : Class<NodeType::Enum> {               \
    using Class<NodeType::Enum>::Class;               \
  };

#define KIVA_EXTERN_TEMPLATE_VARIABLE(Class, Enum, Name) \
  extern template class Class<NodeType::Enum>;           \
  constexpr inline Class<NodeType::Enum> Name;

namespace dsl {

template <NodeType Type>
class FunctionDefBuilder : public DslNodeBuilder {
  FunctionDefBuilder() : DslNodeBuilder(Type) {}

 public:
  // def(name, body)
  FunctionDefBuilder(const char* name, DslNode body);
  // def(name, arguments, returns, body)
  FunctionDefBuilder(const char* name,
                     DslNode arguments,
                     DslNode returns,
                     DslNode body);
  // def(name, arguments, body)
  FunctionDefBuilder(const char* name, DslNode arguments, DslNode body);
};

#define KIVA_FUNCTION_BUILDERS(V)          \
  V(FunctionDefBuilder, kFunctionDef, def) \
  V(FunctionDefBuilder, kAsyncFunctionDef, async_def)

KIVA_FUNCTION_BUILDERS(KIVA_EXTERN_TEMPLATE_ALIAS)

class class_ : public DslNodeBuilder {
  class_() : DslNodeBuilder(NodeType::kClassDef) {}

 public:
  class_(const char* name, DslNode body);
  class_(const char* name, DslNode bases, DslNode body);
};

template <NodeType Type>
class VariableDslNodeBuilder {
 public:
  operator DslNode() const { return Node::NewWithChildren(Type); }
};

template <NodeType Type>
class OptionalValueBuilder : public VariableDslNodeBuilder<Type> {
 public:
  DslNode operator()(DslNode value) const {
    return Node::NewWithChildren(Type, value);
  }
};

#define KIVA_OPTIONAL_VALUE_BUILDERS(V)     \
  V(OptionalValueBuilder, kReturn, return_) \
  V(OptionalValueBuilder, kYield, yield)

KIVA_OPTIONAL_VALUE_BUILDERS(KIVA_EXTERN_TEMPLATE_VARIABLE)

template <NodeType Type>
class ForBuilder : public DslNodeBuilder {
 public:
  ForBuilder(DslNode target, DslNode iter, DslNode body)
      : DslNodeBuilder(Type) {
    AddChildren({target, iter, MakeBody(body)});
  }
};

#define KIVA_FOR_BUILDERS(V) \
  V(ForBuilder, kFor, for_)  \
  V(ForBuilder, kAsyncFor, async_for)

KIVA_FOR_BUILDERS(KIVA_EXTERN_TEMPLATE_ALIAS)

template <NodeType Type>
class ConditionBodyBuilder : public DslNodeBuilder {
 public:
  ConditionBodyBuilder(DslNode condition, DslNode body) : DslNodeBuilder(Type) {
    AddChildren({condition, MakeBody(body)});
  }
};

#define KIVA_CONDITION_BODY_BUILDERS(V)   \
  V(ConditionBodyBuilder, kWhile, while_) \
  V(ConditionBodyBuilder, kIf, if_)       \
  V(ConditionBodyBuilder, kElif, elif)

KIVA_CONDITION_BODY_BUILDERS(KIVA_EXTERN_TEMPLATE_ALIAS)

template <NodeType Type>
class BodyBuilder : public DslNodeBuilder {
  // Prevent things like else_{} and force people to use else_{pass}.
  BodyBuilder();

 public:
  explicit BodyBuilder(DslNode::InitializerList body) : DslNodeBuilder(Type) {
    AddChildren(body);
  }
};

#define KIVA_BODY_BUILDERS(V)       \
  V(BodyBuilder, kElse, else_)      \
  V(BodyBuilder, kTry, try_)        \
  V(BodyBuilder, kFinally, finally) \
  V(BodyBuilder, kDelete, del)

KIVA_BODY_BUILDERS(KIVA_EXTERN_TEMPLATE_ALIAS)

template <NodeType Type>
class WithBuilder : public DslNodeBuilder {
 public:
  WithBuilder(DslNode withitems, DslNode body) : DslNodeBuilder(Type) {
    AddChildren({withitems, MakeBody(body)});
  }
};

#define KIVA_WITH_BUILDERS(V) \
  V(WithBuilder, kWith, with) \
  V(WithBuilder, kAsyncWith, async_with)

KIVA_WITH_BUILDERS(KIVA_EXTERN_TEMPLATE_ALIAS)

class RaiseBuilder : public VariableDslNodeBuilder<NodeType::kRaise> {
 public:
  DslNode operator()(DslNode exc) const;
  DslNode operator()(DslNode exc, DslNode from) const;
};

constexpr inline RaiseBuilder raise;

class except : public DslNodeBuilder {
  except() : DslNodeBuilder(NodeType::kExcept) {}

 public:
  explicit except(DslNode::InitializerList body);
  except(DslNode exc, DslNode body);
  except(DslNode exc, const char* asname, DslNode body);
};

class assert_ : public DslNodeBuilder {
  assert_() : DslNodeBuilder(NodeType::kAssert) {}

 public:
  explicit assert_(DslNode predicate);
  assert_(DslNode predicate, DslNode msg);
};

template <NodeType Type>
class NamesBuilder : public DslNodeBuilder {
 public:
  NamesBuilder(DslNode::InitializerList names);
};

#define KIVA_NAMES_BUILDERS(V)     \
  V(NamesBuilder, kImport, import) \
  V(NamesBuilder, kGlobal, global) \
  V(NamesBuilder, kNonLocal, nonlocal)

KIVA_NAMES_BUILDERS(KIVA_EXTERN_TEMPLATE_ALIAS)

// yield_from(value), await(value)
template <NodeType Type>
class ValueBuilder : public DslNodeBuilder {
 public:
  explicit ValueBuilder(DslNode value);
};

#define KIVA_VALUE_BUILDERS(V)            \
  V(ValueBuilder, kYieldFrom, yield_from) \
  V(ValueBuilder, kAwait, await)          \

KIVA_VALUE_BUILDERS(KIVA_EXTERN_TEMPLATE_ALIAS)

class from : public DslNodeBuilder {
 public:
  explicit from(const char* module);
  DslNode import(DslNode names);
};

class lambda : public DslNodeBuilder {
 public:
  lambda(DslNode args, DslNode expr);
};

#define KIVA_ODDBALL_BUILDERS(V)            \
  V(VariableDslNodeBuilder, kPass, pass)    \
  V(VariableDslNodeBuilder, kBreak, break_) \
  V(VariableDslNodeBuilder, kContinue, continue_)

KIVA_ODDBALL_BUILDERS(KIVA_EXTERN_TEMPLATE_VARIABLE)

#define KIVA_ANNOTATIONS(V)           \
  V(takes, Arguments, arguments)      \
  V(gives, ReturnValue, return_value) \
  V(inherits, Bases, bases)

#define KIVA_DEFINE_ANNOTATION(Name, TParams, Params)   \
  template <typename... TParams>                        \
  DslNode Name(TParams&&... Params) {                   \
    return {DslNode{std::forward<TParams>(Params)}...}; \
  }
KIVA_ANNOTATIONS(KIVA_DEFINE_ANNOTATION)

NameDslNode operator""_n(const char* str, std::size_t len);

// UnaryOp
#define KIVA_DEFINE_OP(Name, Op)                   \
  inline DslNode operator Op(const DslNode& n) {   \
    return Node::NewUnaryOp(NodeType::k##Name, n); \
  }
KIVA_UNARYOP_LIST(KIVA_DEFINE_OP)
#undef KIVA_DEFINE_OP

// BinOp
#define KIVA_DEFINE_OP(Name, Op)                                   \
  inline DslNode operator Op(const DslNode& a, const DslNode& b) { \
    return Node::NewBinOp(a, NodeType::k##Name, b);                \
  }
KIVA_BINOP_LIST(KIVA_DEFINE_OP)
#undef KIVA_DEFINE_OP

// BoolOp
#define KIVA_DEFINE_OP(Name, Op)                                   \
  inline DslNode operator Op(const DslNode& a, const DslNode& b) { \
    return Node::NewBoolOp(a, NodeType::k##Name, b);               \
  }
KIVA_BOOLOP_LIST(KIVA_DEFINE_OP)
#undef KIVA_DEFINE_OP

// Compare
#define KIVA_DEFINE_OP(Name, Op)                                   \
  inline DslNode operator Op(const DslNode& a, const DslNode& b) { \
    return Node::NewCompare(a, NodeType::k##Name, b);              \
  }
KIVA_COMPARE_LIST(KIVA_DEFINE_OP)
#undef KIVA_DEFINE_OP

// AugAssign
#define KIVA_DEFINE_OP(Name, Op)                                   \
  inline DslNode operator Op(const DslNode& a, const DslNode& b) { \
    return Node::NewAugAssign(a, NodeType::k##Name, b);            \
  }
KIVA_AUGASSIGN_LIST(KIVA_DEFINE_OP)
#undef KIVA_DEFINE_OP

#define KIVA_DEFINE_NAME(Name) \
  constexpr inline NameDslNode Name { #Name }

KIVA_DEFINE_NAME(self);
KIVA_DEFINE_NAME(cls);
KIVA_DEFINE_NAME(args);
KIVA_DEFINE_NAME(kwargs);
KIVA_DEFINE_NAME(kwds);
KIVA_DEFINE_NAME(__name__);
KIVA_DEFINE_NAME(__all__);
KIVA_DEFINE_NAME(__file__);
KIVA_DEFINE_NAME(__metaclass__);
KIVA_DEFINE_NAME(__version__);
KIVA_DEFINE_NAME(metaclass);

}  // namespace dsl
}  // namespace kiva
